import json
import os
import sys

import backtrader as bt
import numpy as np
import pandas as pd

from run_backtrader2 import (
    CustomFutureCommission,
    CustomStockCommission,
    MultiFactorStrategy,
    SSPandasData,
)
from utils import get_code_market_data


class LongShortStrategy(MultiFactorStrategy):
    """
    选top_num支股票做多，top_num支股票做空
    """

    params = (
        ("unit_k", 1),  # 成交单位量，
        ("is_stock", False),  # 是否是股票回测
        ("basedir", None),  # 保存回测结果的路径
        ("top_num", 100),  # 每天选取排名靠前的30只股票
        ("max_trade_cent", 0.05),  # 每天每只股票最多成交市场成交量5%
        ("trade_pos", 0.5),  # 最多用现有资金的50%买股票
        ("stock_future", "IF00.IF"),
    )

    def next(self):
        if not hasattr(self, "stock_future_data"):
            self.stock_future_data = next(
                filter(lambda x: x._name == self.p.stock_future, self.datas)
            )

        stock_future_cmi = self.broker.getcommissioninfo(self.stock_future_data)
        top_num = self.p.top_num
        max_trade_cent = self.p.max_trade_cent
        trade_pos = self.p.trade_pos
        now = self.datas[0].datetime.datetime()

        target_value = {
            data: data.pred[0]
            for data in self.datas
            if not np.isnan(data.pred[0])
            if data._name != self.p.stock_future
        }
        # 获取当天每只股票的因子值
        stock_datas = [data for data in self.datas if data._name != self.p.stock_future]
        rank_target = sorted(target_value.items(), key=lambda x: x[1], reverse=True)
        # 对股票因子值从高到低排名

        choose_target = [x[0] for x in rank_target[:top_num]]
        # 选取排名靠前的因子

        print(f"日期：{now} 选择：{sorted([x._name for x in choose_target])}")
        self.every_day_choose.append(
            {
                "datetime": now,
                "choose_stocks": ",".join([x._name for x in choose_target]),
            }
        )
        all_worth = self.broker.getvalue()
        target_pos_mm = {
            data._name: all_worth
            / len(choose_target)
            * trade_pos
            * 1
            / (1 + stock_future_cmi.p.margin)
            // (data.close[0] * self.p.unit_k)
            for data in choose_target
        }

        target_stock_future_pos = -(
            all_worth
            * trade_pos
            * 1
            / (1 + stock_future_cmi.p.margin)
            // (self.stock_future_data.close[0] * 10)
        )
        # 根据当前资产价值计算目标仓位

        cur_stock_future_pos = self.getposition(self.stock_future_data).size // 10

        modify_stock_future_pos = target_stock_future_pos - cur_stock_future_pos

        cur_pos = {
            data._name: self.getposition(data).size // self.p.unit_k
            for data in stock_datas
        }
        modify_pos = {
            data: target_pos_mm.get(data._name, 0) - cur_pos.get(data._name, 0)
            for data in stock_datas
        }
        # 根据当前仓位和目标仓位计算买卖股票数量

        for data, mv in modify_pos.items():
            if mv < 0:
                mv = min(int(data.volume[0] * max_trade_cent), -mv)
                self.sell(data, size=self.p.unit_k * int(mv))
        for data, mv in modify_pos.items():
            if mv > 0:
                mv = min(int(data.volume[0] * max_trade_cent), mv)
                self.buy(data, size=self.p.unit_k * int(mv))

        if modify_stock_future_pos > 0:
            mv = min(
                int(self.stock_future_data.volume[0] * max_trade_cent),
                -modify_stock_future_pos,
            )
            self.buy(self.stock_future_data, size=int(mv) * 10)
        elif modify_stock_future_pos < 0:
            mv = min(
                int(self.stock_future_data.volume[0] * max_trade_cent),
                modify_stock_future_pos,
            )
            self.sell(self.stock_future_data, size=int(mv) * 10)

        self.all_cash.append(
            {
                "datetime": self.datas[0].datetime.datetime(),
                "cash": self.broker.getvalue(),
            }
        )


def main_job(
    result_file: str,
    code_market_list: list,
    start: str,
    end: str,
    cash: float,
    slip: float,
    is_stock: bool,
    cmp_stock_future: str = "IF00.IF",
):
    """
    :param result_file: 保存的模型预测结果
    :param code_market_list: 选择回测的股票代码，为空list则回测所有
    :param start: 回测时间范围“20210101”
    :param end: 结束时间范围“20240701”
    :param cash: 回测使用的金额
    :param slip: 回测滑点，0.0001表示0.01%的滑点
    :param is_stock: 是否回测股票
    :return:
    """
    dir_base = os.path.join(
        os.path.dirname(result_file), f"equity_with_{cmp_stock_future}"
    )
    os.makedirs(dir_base, exist_ok=True)
    df = pd.read_csv(result_file)
    if "volume" not in df.columns:
        df["volume"] = 1000
    if "openInterest" not in df.columns:
        df["openinterest"] = 0.0
    else:
        df.rename(columns={"openInterest": "openinterest"}, inplace=True)
    df.drop("date_time", inplace=True, axis=1)
    df["datetime"] = pd.to_datetime(
        df["time"] + 8 * 60 * 60 * 1000, unit="ms", utc=False
    )  # 时区转换
    df.drop("time", inplace=True, axis=1)
    df.reset_index(drop=True, inplace=True)
    if code_market_list:
        df = df[df["code_market"].isin(code_market_list)].reset_index()
    df = df[
        (df["datetime"] >= pd.to_datetime(start))
        & (df["datetime"] < pd.to_datetime(end))
    ]
    df.set_index("datetime", inplace=True)
    df.sort_index(inplace=True)

    df_if = get_code_market_data(
        cmp_stock_future, "", "", "1d", "front_ratio", fill_data=True
    )
    df_if["datetime"] = pd.to_datetime(
        df_if["time"] + 8 * 60 * 60 * 1000, unit="ms", utc=False
    )  # 时区转换
    df_if.drop("time", inplace=True, axis=1)
    df_if.reset_index(drop=True, inplace=True)
    df_if = df_if[
        (df_if["datetime"] >= pd.to_datetime(start))
        & (df_if["datetime"] < pd.to_datetime(end))
    ]
    df_if.set_index("datetime", inplace=True)
    df_if.sort_index(inplace=True)
    df_if["pred"] = 0

    # 创建一个Cerebro引擎
    cerebro = bt.Cerebro()

    # 添加策略
    cerebro.addstrategy(
        LongShortStrategy,
        unit_k=100 if is_stock else 10,
        is_stock=is_stock,
        basedir=dir_base,
        stock_future=cmp_stock_future,
    )
    dt_idx = df.index.drop_duplicates()
    stock_data = {}
    for xcode, df_group in df.groupby("code_market"):
        # 添加数据
        df_new = df_group.reindex(dt_idx).ffill().bfill()
        df_new.sort_index(inplace=True)
        stock_data[xcode] = df_new
    for xcode, df_group in stock_data.items():
        data_feed = SSPandasData(dataname=df_group.sort_index(), name=xcode)
        cerebro.adddata(data_feed)

    data_if_feed = SSPandasData(dataname=df_if.sort_index(), name=xcode)
    cerebro.adddata(data_if_feed, name=cmp_stock_future)

    cerebro.broker.setcash(cash)
    cmi = CustomStockCommission() if is_stock else CustomFutureCommission()
    cerebro.broker.addcommissioninfo(cmi)
    cerebro.broker.addcommissioninfo(CustomFutureCommission(), name=cmp_stock_future)

    cerebro.broker.set_slippage_perc(perc=slip)
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name="sharpe_ratio")
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name="drawdown")
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name="annual_return")

    start_cash = cerebro.broker.getvalue()
    results = cerebro.run()
    end_cash = cerebro.broker.getvalue()
    sharpe_ratio = results[0].analyzers.sharpe_ratio.get_analysis()
    max_drawdown = results[0].analyzers.drawdown.get_analysis()
    annual_return = results[0].analyzers.annual_return.get_analysis()

    summary = dict(
        result_file=result_file,
        start=start,
        end=end,
        cash=cash,
        start_cash=start_cash,
        end_cash=end_cash,
        sharpe_ratio=sharpe_ratio,
        max_drawdown=max_drawdown,
        annual_return=annual_return,
        cmi=str(cmi.p),
        code_market_list=code_market_list,
    )
    print(summary)
    with open(os.path.join(dir_base, "backtest.json"), "w") as f:
        json.dump(summary, f, ensure_ascii=False, indent=2)

    print("Backtest result saved as backtest_result.png")


if __name__ == "__main__":
    datafile, start, end = sys.argv[1:]
    main_job(
        datafile,
        # ["601828.SH","603789.SH","601229.SH","600606.SH","688535.SH","603377.SH","601528.SH","601916.SH","603956.SH","603073.SH","603117.SH","688319.SH"],
        [],
        start,
        end,
        10000000,
        slip=0.0002,
        is_stock=True,
    )
